'''
Created on 2012-4-23

@author: nzhao
'''
import numpy as np

import materials 
from Spin import Spin

class MultiSpin():
    '''
    a class of Multi-spins
    '''


    def __init__(self, spinlist=[], cSpin=[], x=np.array([1., 0., 0.]), y=np.array([0., 1., 0.]), z=np.array([0., 0., 1.])):
        '''
        Constructor
        '''
        self.spinNum=len(spinlist)
        self.spinList=spinlist
        self.interaction=[]

        self.centralSpinNum=len(cSpin)
        self.centralSpin=cSpin        
        self.couping=[]
        
        self.xDirection=x
        self.yDirection=y
        self.zDirection=z
           
    def add_central_spin(self, spin):
        self.centralSpin.append(spin)
        self.centralSpinNum=len(self.centralSpin)
    
    def add_spin(self, spin):
        self.spinList.append(spin)
        self.spinNum = len(self.spinList)
        
    def rm_spin(self, idx):
        self.spinList.pop(idx)
        self.spinNum = len(self.spinList)
    
    def calculate_interaction(self):
        self.interaction=np.array( \
                        [self.dipolar(self.spinList[i], self.spinList[j]) \
                          for i in range(self.spinNum) \
                          for j in range(self.spinNum) \
                          if i > j ])
    def calculate_coupling(self, isDipolar=True):
        if isDipolar:
            self.couping=np.array(\
                        [self.dipolar(spin_i, spin_j) \
                          for spin_i in self.centralSpin\
                          for spin_j in self.spinList ])
        else:
            pass
    
    def index2(self, i, j):
        return i*(i-1)/2+j
    
    def dipolar(self, spin_i, spin_j):
        gi = spin_i.spinRelGamma
        gj = spin_j.spinRelGamma
        d = gi * gj * materials.SpinSpinInteraction
        
        ri = spin_i.spinPosition
        rj = spin_j.spinPosition
        rij = ri - rj
        r_norm = np.linalg.norm(rij)
        nij = rij / r_norm 
         
        nx, ny, nz = np.dot(nij, self.xDirection), np.dot(nij, self.yDirection), np.dot(nij, self.zDirection)
        
        res = d / r_norm**3 * np.array([1.0 - 3.0 * nx*nx, 1.0 - 3.0 * ny*ny, 1.0 - 3.0 * nz*nz,\
                   -3.0 * nx*ny, -3.0 * ny*nz, -3.0 * nz*nx])
        return res
        
        
if __name__ == '__main__':
    multiSpin1=MultiSpin()
    s1=Spin()
    multiSpin1.add_spin(s1)
    print multiSpin1.spinNum
    